import numpy as np
from . import LocalSearchOperator, register_operator


class BlockInsertionOperator(LocalSearchOperator):
    """块插入(Block Insertion)局部搜索算子"""

    def __init__(self):
        self.name = "block_insertion"
        self.supported_problems = ["tsp"]
        self.max_block_size = 5  # 最大块大小

    def generate_neighborhood(self, solution, problem_instance):
        """生成块插入邻域"""
        neighborhood = []
        n = len(solution)

        # 尝试不同的块大小
        for block_size in range(2, min(self.max_block_size + 1, n // 2)):
            # 对每个可能的块起始位置
            for i in range(0, n - block_size + 1):
                # 对每个可能的插入位置
                for j in range(0, n - block_size + 1):
                    if i == j or (j >= i and j < i + block_size) or (i >= j and i < j + block_size):
                        continue  # 跳过重叠位置

                    # 创建新解：将块从i到i+block_size-1插入到位置j
                    block = solution[i:i + block_size]

                    # 移除块
                    without_block = np.concatenate([solution[:i], solution[i + block_size:]])

                    # 插入到位置j
                    if j > i:
                        j_adjusted = j - block_size  # 因为移除了块，索引需要调整
                    else:
                        j_adjusted = j

                    new_solution = np.concatenate([
                        without_block[:j_adjusted],
                        block,
                        without_block[j_adjusted:]
                    ])

                    neighborhood.append(new_solution)

        return neighborhood

    def efficient_block_insertion(self, solution, problem_instance, intensity=10):
        """高效的块插入实现，直接计算成本变化"""
        n = len(solution)
        current_solution = solution.copy()
        current_cost = problem_instance.evaluate(current_solution)
        improved = True
        iterations = 0

        while improved and iterations < intensity:
            improved = False

            # 尝试多个随机块插入操作
            for _ in range(min(50, n)):  # 限制尝试次数
                # 随机选择块大小、起始位置和插入位置
                block_size = np.random.randint(2, min(self.max_block_size + 1, n // 2))
                i = np.random.randint(0, n - block_size + 1)  # 块起始位置
                j = np.random.randint(0, n - block_size + 1)  # 插入位置

                # 确保位置不重叠
                if i == j or (j >= i and j < i + block_size) or (i >= j and i < j + block_size):
                    continue

                # 计算块插入前的相关边成本
                # 块前后的边
                prev_i = (i - 1) % n
                next_i = (i + block_size) % n

                # 插入位置前后的边
                if j == 0:
                    prev_j = n - 1
                    next_j = 0
                else:
                    prev_j = j - 1
                    next_j = j

                # 计算插入前的成本
                old_cost = (
                        problem_instance.distance_matrix[current_solution[prev_i]][current_solution[i]] +
                        problem_instance.distance_matrix[current_solution[i + block_size - 1]][
                            current_solution[next_i]] +
                        problem_instance.distance_matrix[current_solution[prev_j]][current_solution[j]]
                )

                # 执行块插入操作
                block = current_solution[i:i + block_size]
                without_block = np.concatenate([
                    current_solution[:i],
                    current_solution[i + block_size:]
                ])

                # 调整插入位置索引
                if j > i:
                    j_adjusted = j - block_size
                else:
                    j_adjusted = j

                new_solution = np.concatenate([
                    without_block[:j_adjusted],
                    block,
                    without_block[j_adjusted:]
                ])

                # 计算插入后的相关边成本
                new_n = n  # 长度不变

                # 确定新位置的前后节点
                if j_adjusted == 0:
                    new_prev = new_n - 1
                    new_next = block_size
                else:
                    new_prev = j_adjusted - 1
                    new_next = j_adjusted + block_size

                # 计算插入后的成本
                new_cost = (
                        problem_instance.distance_matrix[new_solution[new_prev]][new_solution[j_adjusted]] +
                        problem_instance.distance_matrix[new_solution[j_adjusted + block_size - 1]][
                            new_solution[new_next]] +
                        problem_instance.distance_matrix[new_solution[new_prev]][new_solution[new_next]]
                )

                # 计算成本变化（近似值）
                cost_change = new_cost - old_cost

                # 如果成本降低，则接受插入
                if cost_change < 0:
                    current_solution = new_solution
                    current_cost = problem_instance.evaluate(current_solution)  # 直接评估新解
                    improved = True
                    break  # 找到一个改进就跳出内层循环

            iterations += 1

        improvement = problem_instance.evaluate(solution) - current_cost
        return current_solution, improvement

    def __call__(self, solution, problem_instance, intensity=10):
        """应用块插入局部搜索"""
        if not self.is_applicable(problem_instance):
            return solution, 0

        n = len(solution)
        if n < 6:  # 块插入需要至少6个城市（块大小至少为2）
            return solution, 0

        # 使用高效实现
        return self.efficient_block_insertion(solution, problem_instance, intensity)


# 注册算子
register_operator("block_insertion", BlockInsertionOperator())